da260d
@@ -87,6 +87,8 @@
 
   private final Object endOfEvents = new Object();
 
+  private int totalEventCount = 0;
+
   public DynamicPartitionPruner() {
   }
 
@@ -114,7 +116,7 @@
public void prune(MapWork work, JobConf jobConf, InputInitializerContext context
     // synchronous event processing loop. Won't return until all events have
     // been processed.
     this.processEvents();
-    this.prunePartitions(work);
+    this.prunePartitions(work, context);
     LOG.info("Ok to proceed.");
   }
 
@@ -163,12 +165,22 @@
public void initialize(MapWork work, JobConf jobConf) throws SerDeException {
     }
   }
 
-  private void prunePartitions(MapWork work) throws HiveException {
+  private void prunePartitions(MapWork work, InputInitializerContext context) throws HiveException {
+    int expectedEvents = 0;
     for (String source : this.sourceInfoMap.keySet()) {
       for (SourceInfo si : this.sourceInfoMap.get(source)) {
+        int taskNum = context.getVertexNumTasks(source);
+        LOG.info("Expecting " + taskNum + " events for vertex " + source);
+        expectedEvents += taskNum;
         prunePartitionSingleSource(source, si, work);
       }
     }
+
+    // sanity check. all tasks must submit events for us to succeed.
+    if (expectedEvents != totalEventCount) {
+      LOG.error("Expecting: " + expectedEvents + ", received: " + totalEventCount);
+      throw new HiveException("Incorrect event count in dynamic parition pruning");
+    }
   }
 
   private void prunePartitionSingleSource(String source, SourceInfo si, MapWork work)
@@ -396,7 +408,8 @@
public int read(byte[] bytes, int off, int len) throws IOException {
   public void addEvent(InputInitializerEvent event) {
     synchronized(sourcesWaitingForEvents) {
       if (sourcesWaitingForEvents.contains(event.getSourceVertexName())) {
-          queue.offer(event);
+        ++totalEventCount;
+        queue.offer(event);
       }
     }
   }
